package easy.trace.core.context;

import com.alibaba.ttl.TransmittableThreadLocal;
import easy.trace.common.Annotation;
import easy.trace.common.BinaryAnnotation;
import easy.trace.common.Endpoint;
import easy.trace.common.Span;
import easy.trace.core.common.Constants;
import easy.trace.core.dto.*;
import easy.trace.core.util.IdGen;

import java.util.Optional;

/**
 * @author easy
 */
public class SpanContext {

    private final IdGen idGen;
    private ThreadLocal<Span> spanThreadLocal;
    private ThreadLocal<Span> currentSpan = new ThreadLocal<>();
    private volatile boolean notFoundTransmittableThreadLocalClass = false;

    public SpanContext() {
        this.idGen = IdGen.get();
        try {
            this.spanThreadLocal = new TransmittableThreadLocal<>();
        } catch (Throwable e) {
            notFoundTransmittableThreadLocalClass = true;
            System.out.println(e.getLocalizedMessage());
        } finally {
            if (notFoundTransmittableThreadLocalClass) {
                this.spanThreadLocal = new ThreadLocal<>();
            }
        }
    }

    public void removeCurrentSpan() {
        currentSpan.remove();
    }

    public Span getCurrentSpan() {
        return currentSpan.get();
    }

    public void setCurrentSpan(Span span) {
        currentSpan.set(span);
    }

    public void removeParentSpan() {
        spanThreadLocal.remove();
    }

    public Span getParentSpan() {
        return spanThreadLocal.get();
    }

    public void setParentSpan(Span span) {
        spanThreadLocal.set(span);
    }

    public Span createNewSpan(String serviceId, String methodName, String eventType) {
        Long traceId = Optional.ofNullable(this.genTraceId()).orElse(0L);
        Long id = Optional.ofNullable(this.genId()).orElse(0L);
        return buildSimpleSpan(serviceId, methodName, eventType, traceId, 0L, id);
    }

    public Span createNewSpan(String serviceId, String methodName, String eventType, Long traceId, Long parentId, Long id) {
        return buildSimpleSpan(serviceId, methodName, eventType, traceId, parentId, id);
    }

    public Span createNewSpan(String serviceId, String methodName, String eventType, Long traceId, Long parentId) {
        Long id = Optional.ofNullable(this.genId()).orElse(0L);
        return buildSimpleSpan(serviceId, methodName, eventType, traceId, parentId, id);
    }

    public Annotation createNewAnnotation(String value, Long timestamp) {
        return createNewAnnotation(value, timestamp, "127.0.0.1", 0);
    }

    public Annotation createNewAnnotation(String value, String host, int port) {
        return createNewAnnotation(value, System.currentTimeMillis(), host, port);
    }

    public Annotation createNewAnnotation(String value) {
        return createNewAnnotation(value, System.currentTimeMillis(), "127.0.0.1", 0);
    }

    public BinaryAnnotation createNewTimeOutAnnotation(String message, String messageType) {
        return createNewBinaryAnnotation(messageType, message, ExceptionType.TIMEOUT, 0);
    }

    public BinaryAnnotation createNewTimeOutAnnotation(String message) {
        return createNewBinaryAnnotation(Constants.DUBBO_EXCEPTION, message, ExceptionType.TIMEOUT, 0);
    }

    private static BinaryAnnotation createNewBinaryAnnotation(String key, String message, ExceptionType type, int execTime) {
        BinaryAnnotation exAnnotation = new BinaryAnnotation();
        exAnnotation.setKey(key);
        exAnnotation.setValue(message);
        exAnnotation.setType(type.name());
        exAnnotation.setDuration(execTime);
        exAnnotation.setHost(new Endpoint());
        return exAnnotation;
    }

    private static Annotation createNewAnnotation(String value, Long timestamp, String ip, int port) {
        Annotation annotation = new Annotation();
        annotation.setValue(value);
        annotation.setTimestamp(timestamp);
        annotation.setHost(new Endpoint(ip, port));
        return annotation;
    }

    private static Span buildSimpleSpan(String serviceId, String methodName, String eventType, Long traceId, Long parentId, Long id) {
        Span span = new Span();
        span.setServiceId(serviceId);
        span.setSpanName(methodName);
        span.setSample(true);
        span.setTraceId(traceId);
        span.setParentId(parentId);
        span.setId(id);
        span.setEventData(eventType);
        return span;
    }

    public Long genTraceId() {
        return Math.abs(idGen.nextId());
    }

    public Long genId() {
        return idGen.nextId();
    }

}
